package tuple.data;

import org.yaml.snakeyaml.Yaml;
import polyglot.ast.New;
import soot.Body;
import soot.SootMethod;
import soot.Value;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class NewDataRule {

    public static final Map<String, Pattern> DATA_PATTERNS = new HashMap<>();
    private static  String pathToDataOntology =
            "C:\\Users\\Yalin Feng\\Code\\AndroidPrivacy\\Project\\data\\ontology\\data_ontology.yml";
//            "/home/fyl/project/code/AndroPrivacyConsistency/data/ontology/data_ontology.yml";

    private static boolean isLoaded = false;

    public static void ensureDataLoaded() throws IOException {
        if (!isLoaded) {
            loadDataFromYaml();
            isLoaded = true;
        }
    }

    public static void  setDataOntologyPath(String pathToDataOntology){
        //if path points to a real file
        if (Files.exists(Paths.get(pathToDataOntology))) {
            NewDataRule.pathToDataOntology = pathToDataOntology;
        }
    }
    private static void loadDataFromYaml() throws IOException {
        try (InputStream inputStream = Files.newInputStream(Paths.get(pathToDataOntology))) {
            Yaml yaml = new Yaml();
            Map<String, Object> loadedYaml = yaml.load(inputStream);
            NewDataRule.DATA_PATTERNS.putAll(parseLoaded(loadedYaml));
        }
    }

    private static Map<String, Pattern> parseLoaded(Map<String, Object> loadedYaml) {
        Map<String, Pattern> patterns = new HashMap<>();
        Queue<Map<String, Object>> queue = new ArrayDeque<>();
        queue.add(loadedYaml);
        while (!queue.isEmpty()) {
            Map<String, Object> cur = queue.poll();
            if (cur.containsKey("children")) {
                queue.addAll((Collection<? extends Map<String, Object>>) cur.get("children"));
            }
            String name = (String) cur.get("name");
            if (cur.get("patterns") instanceof ArrayList) {
                ArrayList list = (ArrayList) cur.get("patterns");
                StringBuilder allPtn = null;
                for (Object o : list) {
                    if (o instanceof String s) {
                        if (allPtn == null)
                            allPtn = new StringBuilder(s);
                        else {
                            allPtn.append("|").append(s);
                        }
                    }
                }
                if (allPtn != null) {
                    patterns.put(name, Pattern.compile(allPtn.toString()));
                }
            } else if (cur.get("patterns") instanceof String ptns) {
                Pattern p = Pattern.compile(ptns);
                patterns.put(name, p);
            }

        }
        return patterns;
    }

    public static Set<DataKind> checkClassNameAsData(@Nonnull Set<String> splitName) {
        Set<String> strings = DATA_PATTERNS.keySet();
        Set<DataKind> dataKinds = new HashSet<>();
        // 这三个是不可能在包名里的; the three data items are not possible in the package name
        strings.remove("COARSE_LOCATION");
        strings.remove("FINE_LOCATION");
        strings.remove("CONSTANT_STRING");
        // 只靠类名判断数据不可靠,加入到 potentialDataSet
        for (String key : strings) {
            Pattern pattern = DATA_PATTERNS.get(key);
            splitName.forEach(s -> {
                Matcher matcher = pattern.matcher(s.toLowerCase());
                if (matcher.find()) {
                    dataKinds.add(DataKind.valueOf(key.toUpperCase()));
                }
            });
        }
        return dataKinds;
    }

    public static Set<DataKind> checkMethodNameAsData(@Nonnull String name) {
        Set<String> strings = DATA_PATTERNS.keySet();
        Set<DataKind> dataKinds = new HashSet<>();
        strings.remove("CONSTANT_STRING");
        // 只靠类名判断数据不可靠,加入到 potentialDataSet
        for (String key : strings) {
            Pattern pattern = DATA_PATTERNS.get(key);
            Matcher matcher = pattern.matcher(name.toLowerCase());
            if (matcher.find()) {
                dataKinds.add(DataKind.valueOf(key.toUpperCase()));
            }
        }
        return dataKinds;
    }

    //"advertis(?!id\b)"
    public static Set<DataKind> checkMethodBodyAsData(@Nonnull Set<Value> noDupValues) {
        Set<String> strings = DATA_PATTERNS.keySet();
        Set<DataKind> dataKinds = new HashSet<>();
        for (String key : strings) {
            Pattern pattern = DATA_PATTERNS.get(key);
            noDupValues.forEach(value -> {
                Matcher matcher = pattern.matcher(value.toString().toLowerCase());
                if (matcher.find()) {
                    dataKinds.add(DataKind.valueOf(key.toUpperCase()));
                }
            });
        }

        return dataKinds;
    }

    public static Set<DataKind> checkMethodData(SootMethod sm) {
        Set<DataKind> data = new HashSet<>();
        // 1. className
        String className = sm.getDeclaringClass().getName();
        Set<String> splitName = Arrays.stream(className.split("\\."))
                .filter(x -> x.length() > 1)
                .collect(Collectors.toSet());
        data.addAll(NewDataRule.checkClassNameAsData(splitName));

        // 2. methodName
        data.addAll(NewDataRule.checkMethodNameAsData(sm.getSignature()));

        try {
            // 3. methodBody
            if (sm.hasActiveBody()) {
                Body body = sm.retrieveActiveBody();
                Set<Value> noDupValues = body.getUnits().stream()
                        .flatMap(x -> x.getUseAndDefBoxes().stream())
                        .map(x -> x.getValue())
                        .collect(Collectors.toSet());//JimpleLocal, JSpecialInvokeExpre,JNewExpr,JInstanceFieldRef,ThisRef
                data.addAll(NewDataRule.checkMethodBodyAsData(noDupValues));
            }
        } catch (RuntimeException e) {
            e.printStackTrace();
        }
        return data;
    }
}
